#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 接收 /force_data 话题
# 如果接收到的力不超过期望力020，则控制机械臂往初始位置运动

import time
import numpy as np
import rospy
from moveit_commander import MoveGroupCommander
from geometry_msgs.msg import Pose, TwistStamped

class ForceControl:
    def __init__(self):
        self.force_torque_filtered = np.zeros((3, 1))
        self.force_torque_reference = np.zeros((3, 1))
        self.force_torque_reference[1] = 0  # 参考力
        self.position_reference = np.zeros((3, 1))
        self.tcp_frame_cur_p = np.zeros((3, 1))
        self.tcp_frame_cur_M = np.eye(3)
        self.velocity = np.zeros((3, 1))
        self.acceleration = np.zeros((3, 1))
        self.K_position = np.eye(3) * 1.1
        self.B_position = np.eye(3) * 0.01
        self.M_position = np.eye(3) * 280
        self.dt = 1
        self.first_get_coord = True
        self.group = MoveGroupCommander("manipulator")

        # 订阅 /force_data 话题
        self.wrench_sub = rospy.Subscriber('/force_data', TwistStamped, self.force_callback)
        
    def force_callback(self, msg):
        current_force = msg.twist.linear.z
        if (self.force_torque_filtered[0] > 0 and current_force < 0) or (self.force_torque_filtered[0] < 0 and current_force > 0):
            self.first_get_coord = True # 当正负变化时重置坐标获取标志
        self.force_torque_filtered[0] = msg.twist.linear.z*5
        self.force_torque_filtered[1] = 0
        self.force_torque_filtered[2] = 0

    def update_current_pose(self):
        current_pose = self.group.get_current_pose().pose
        self.tcp_frame_cur_p[0, 0] = current_pose.position.x
        self.tcp_frame_cur_p[1, 0] = current_pose.position.y
        self.tcp_frame_cur_p[2, 0] = current_pose.position.z
        if self.first_get_coord:
            self.position_reference[0, 0] = current_pose.position.x
            self.position_reference[1, 0] = current_pose.position.y
            self.position_reference[2, 0] = current_pose.position.z
            self.first_get_coord = False

    def calculate_new_pose(self):
        force_difference = self.force_torque_filtered - self.force_torque_reference
        position_difference = self.tcp_frame_cur_p - self.position_reference
        Kp_position_term = np.dot(self.K_position, position_difference)
        B_velocity_term = np.dot(self.B_position, self.velocity)
        
        self.acceleration = np.dot(np.linalg.inv(self.M_position), 
                                   (force_difference - Kp_position_term - B_velocity_term))
        self.velocity += (self.acceleration * self.dt)
        delta_position = self.velocity * self.dt
        self.velocity = 0.001
        
        new_position = self.tcp_frame_cur_p - delta_position
        print("新位置:\n", new_position)
        return new_position
    
    def force_below_threshold(self, threshold, tolerance=0.1):
        return np.all(np.abs(self.force_torque_filtered) <= threshold + tolerance)

def main():
    rospy.init_node('force_control_node')  # 初始化ROS节点
    group_name = "manipulator"
    move_group = MoveGroupCommander(group_name)
    force_control = ForceControl()
    
    # 设置初始位置
    move_group.set_named_target('middle')
    move_group.go()
    move_group.stop()
    move_group.clear_pose_targets()
    
    try:
        while not rospy.is_shutdown():
            force_control.update_current_pose()
            
            # 检查接收到的力是否超过期望力
            if force_control.force_torque_filtered[0] == 0:
                move_group.set_named_target('middle')
            else:
                new_position = force_control.calculate_new_pose()
                
                waypoints = []
                pose_target = move_group.get_current_pose().pose
                for i in range(10):
                    pose_target.position.x += (new_position[0, 0] - pose_target.position.x) / 10
                    pose_target.position.y = 0
                    pose_target.position.z += (new_position[2, 0] - pose_target.position.z) / 10
                    waypoints.append(pose_target)
                
                # 计算路径并执行
                (plan, fraction) = move_group.compute_cartesian_path(waypoints, eef_step=0.01, jump_threshold=0.0)
                move_group.execute(plan, wait=True)

                # 检查 z 和 y 值是否超出限制
                if pose_target.position.z > 0.35 or pose_target.position.y > 0.001:
                    if pose_target.position.z > 0.35:
                        print("Z value exceeds 0.35. Exiting the process.")
                    if pose_target.position.y > 0.001:
                        print("Y value exceeds 0.001. Exiting the process.")
                    rospy.sleep(2)
                    move_group.set_named_target('middle')
                    move_group.go(wait=True)
                    break  # 结束循环和进程
            
            time.sleep(1)
    except KeyboardInterrupt:
        print("Motion stopped by user")
    finally:
        move_group.stop()
        move_group.clear_pose_targets()

if __name__ == "__main__":
    main()